import os
import sys
import numpy as np
from mrcnn import visualize
import mdf as coco
ROOT_DIR = os.path.abspath("./")
sys.path.append(ROOT_DIR)
config = coco.CocoConfig()
COCO_DIR = "./coco_mdf"
# Load dataset
dataset = coco.CocoDataset()
dataset.load_coco(COCO_DIR, "train")
dataset.prepare()

print("Image Count: {}".format(len(dataset.image_ids)))
print("Class Count: {}".format(dataset.num_classes))
for i, info in enumerate(dataset.class_info):
    print("{:3}. {:50}".format(i, info['name']))
# Load and display random samples
image_ids = np.random.choice(dataset.image_ids, 20)
print(image_ids)
for image_id in image_ids:
    image = dataset.load_image(image_id)
    mask, class_ids = dataset.load_mask(image_id)
    print(mask.shape)
    visualize.display_top_masks(image, mask, class_ids, dataset.class_names)